package com.ajocer.springbootinit.rag;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * 简化版Milvus向量存储实现
 * 用于解决兼容性问题
 */
@Slf4j
@Component
public class SimpleMilvusVectorStore implements EmbeddingStore<TextSegment> {

    // 存储嵌入向量
    private final Map<String, Embedding> embeddings = new ConcurrentHashMap<>();
    
    // 存储文本段
    private final Map<String, TextSegment> segments = new ConcurrentHashMap<>();
    
    // ID生成器
    private final AtomicLong idGenerator = new AtomicLong();

    @Override
    public void add(String id, Embedding embedding) {
        embeddings.put(id, embedding);
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings) {
        List<String> ids = new ArrayList<>(embeddings.size());
        for (Embedding embedding : embeddings) {
            ids.add(add(embedding));
        }
        return ids;
    }

    @Override
    public String add(Embedding embedding) {
        String id = String.valueOf(idGenerator.incrementAndGet());
        embeddings.put(id, embedding);
        return id;
    }

    @Override
    public String add(Embedding embedding, TextSegment textSegment) {
        String id = String.valueOf(idGenerator.incrementAndGet());
        embeddings.put(id, embedding);
        segments.put(id, textSegment);
        return id;
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings, List<TextSegment> textSegments) {
        if (embeddings.size() != textSegments.size()) {
            throw new IllegalArgumentException("嵌入向量和文本段数量不匹配");
        }
        
        List<String> ids = new ArrayList<>(embeddings.size());
        for (int i = 0; i < embeddings.size(); i++) {
            ids.add(add(embeddings.get(i), textSegments.get(i)));
        }
        return ids;
    }

    @Override
    public List<EmbeddingMatch<TextSegment>> findRelevant(Embedding referenceEmbedding, int maxResults, double minScore) {
        // 计算所有嵌入向量与参考向量的相似度
        List<EmbeddingMatch<TextSegment>> matches = new ArrayList<>();
        
        for (Map.Entry<String, Embedding> entry : embeddings.entrySet()) {
            String id = entry.getKey();
            Embedding embedding = entry.getValue();
            
            // 计算余弦相似度
            double similarity = cosineSimilarity(referenceEmbedding, embedding);
            
            // 过滤低于最小分数的结果
            if (similarity >= minScore) {
                TextSegment segment = segments.get(id);
                if (segment != null) {
                    // 使用工厂类创建EmbeddingMatch实例
                    EmbeddingMatch<TextSegment> match = EmbeddingMatchFactory.create(similarity, segment);
                    matches.add(match);
                }
            }
        }
        
        // 按相似度降序排序
        matches.sort((a, b) -> Double.compare(b.score(), a.score()));
        
        // 限制结果数量
        return matches.stream()
                .limit(maxResults)
                .collect(Collectors.toList());
    }

    /**
     * 计算余弦相似度
     */
    private double cosineSimilarity(Embedding e1, Embedding e2) {
        List<Float> v1 = e1.vectorAsList();
        List<Float> v2 = e2.vectorAsList();
        
        if (v1.size() != v2.size()) {
            throw new IllegalArgumentException("向量维度不匹配");
        }
        
        double dotProduct = 0.0;
        double norm1 = 0.0;
        double norm2 = 0.0;
        
        for (int i = 0; i < v1.size(); i++) {
            dotProduct += v1.get(i) * v2.get(i);
            norm1 += Math.pow(v1.get(i), 2);
            norm2 += Math.pow(v2.get(i), 2);
        }
        
        // 避免除以零
        if (norm1 == 0 || norm2 == 0) {
            return 0.0;
        }
        
        return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
    }
} 